#!/usr/bin/python
# -*-coding:utf-8-*-
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns

import os
import numpy as np
import pandas as pd

import scipy.cluster.hierarchy as sch

# from sklearn.cluster.bicluster import SpectralCoclustering
from sklearn.cluster import SpectralCoclustering

from sklearn.manifold import MDS
from sklearn.cluster import KMeans

# from dir_info import cache_data_dir

period_corr_data_save_dir = './zbc_gplearn_factor_mining/corr_analysis'
period_corr_plot_save_dir = './zbc_gplearn_factor_mining/corr_analysis/plot'

if not os.path.exists(period_corr_plot_save_dir):
    os.makedirs(period_corr_plot_save_dir)

mean_corr_data_save_filename = 'zbc_gp_ub_factor_mean_corr_data_test_period_1'

prefix_labl = 'zbc_gp_ub_factor_mean_corr_data_test_period_1'

# 画图参数
# figsize = (15, 10)
figsize = (20, 20)
# labelsize = None
labelsize = 5

###########################################################
def plot(df, save_path, figsize=figsize):
    corr = df.copy()

    fig, ax = plt.subplots(figsize=figsize)

    # cax = ax.matshow(corr, cmap='RdYlGn')
    sns.heatmap(corr,
                vmax=1.0,
                vmin=-1.0,
                square=True,
                # center=0.0,
                ax=ax,
                # linewidths=0.005,
                cmap='RdYlGn')

    plt.xticks(range(len(corr.columns)), corr.columns, rotation=90)
    plt.yticks(range(len(corr.columns)), corr.columns)

    # fig.colorbar(cax, ticks=[-1, 0, 1], aspect=40, shrink=0.8)

    plt.tight_layout()

    plt.savefig(save_path, dpi=200)
    plt.close('all')

def plot_abs(df, save_path, figsize=figsize, labelsize=labelsize):
    corr = df.copy()

    fig, ax = plt.subplots(figsize=figsize)

    # cax = ax.matshow(corr, cmap='RdYlGn')
    sns.heatmap(corr,
                vmin=0.0,
                vmax=1.0,
                square=True,
                ax=ax,
                # linewidths=0.005,
                cmap='RdYlGn')

    plt.xticks(range(len(corr.columns)), corr.columns, rotation=90)
    plt.yticks(range(len(corr.columns)), corr.columns)

    if labelsize is not None:
        plt.tick_params(labelsize=labelsize)

    # fig.colorbar(cax, ticks=[0, 1], aspect=40, shrink=0.8)

    plt.tight_layout()

    plt.savefig(save_path, dpi=200)
    plt.close('all')

###
### TODO - read data
try:
    mean_period_corr_data = pd.read_hdf(os.path.join(period_corr_data_save_dir,
                                                     mean_corr_data_save_filename + '.h5'))
except:
    mean_period_corr_data = pd.read_excel(os.path.join(period_corr_data_save_dir,
                                                     mean_corr_data_save_filename + '.xlsx'))
    mean_period_corr_data = mean_period_corr_data.set_index('factor_name')

# TODO - 处理nan数据
# TODO - 全空，则去掉
factor_nan_num = mean_period_corr_data.isnull().sum(axis=1)

factor_all_num = factor_nan_num.shape[0]

factor_nan_num = factor_nan_num[factor_nan_num <= 0.5*factor_all_num]

selected_factor_name = factor_nan_num.index

analysis_mean_period_corr_data = mean_period_corr_data.loc[selected_factor_name, selected_factor_name].copy()

# 其余全设为0.0
analysis_mean_period_corr_data = analysis_mean_period_corr_data.fillna(0.0)

analysis_mean_period_corr_data = analysis_mean_period_corr_data.abs()

# 参数
n_cluster = 5
figsize = (20, 20)

### TODO - 相关性分析
save_path = os.path.join(period_corr_plot_save_dir, prefix_labl+'_processed_corr_data_map.png')
plot_abs(analysis_mean_period_corr_data, save_path=save_path, figsize=figsize)

# TODO - SpectralCoclustering
scc = SpectralCoclustering(n_clusters=n_cluster, random_state=123)
scc.fit(analysis_mean_period_corr_data)

corr_scc_data = analysis_mean_period_corr_data.iloc[np.argsort(scc.row_labels_)]
corr_scc_data = corr_scc_data.iloc[:, np.argsort(scc.row_labels_)]
corr_scc_data.to_excel(os.path.join(period_corr_data_save_dir, prefix_labl+'_spectral_coclustering_corr_data_map.xlsx'))

save_path = os.path.join(period_corr_plot_save_dir, prefix_labl+'_spectral_coclustering_corr_data_map.png')
# plot(corr_scc_data, save_path=save_path, figsize=figsize)
plot_abs(corr_scc_data, save_path=save_path, figsize=figsize)

# TODO - linkage and fcluster
X = analysis_mean_period_corr_data.values

d = sch.distance.pdist(X)

L = sch.linkage(d, method='complete')

ind = sch.fcluster(L, 0.5*d.max(), 'distance', depth=5)

index_columns = [analysis_mean_period_corr_data.columns.tolist()[i] for i in list((np.argsort(ind)))]

corr_linkage_fc_data = analysis_mean_period_corr_data.loc[index_columns, index_columns]
corr_linkage_fc_data.to_excel(os.path.join(period_corr_data_save_dir, prefix_labl+'_linkage_fcluster_corr_data_map.xlsx'))

save_path = os.path.join(period_corr_plot_save_dir, prefix_labl+'_linkage_fcluster_corr_data_map.png')
# plot(corr_linkage_fc_data, save_path=save_path, figsize=figsize)
plot_abs(corr_linkage_fc_data, save_path=save_path, figsize=figsize)

# TODO - Kmeans + MDS
embedding = MDS(n_components=5)

analysis_mean_period_corr_data_transformed = embedding.fit_transform(analysis_mean_period_corr_data.values)

kmeans = KMeans(n_clusters=n_cluster, algorithm='full')
kmeans.fit(analysis_mean_period_corr_data_transformed)

kmeans_ind = kmeans.predict(analysis_mean_period_corr_data_transformed)

corr_mds_kmeans_data = analysis_mean_period_corr_data.iloc[np.argsort(kmeans_ind)]
corr_mds_kmeans_data = corr_mds_kmeans_data.iloc[:, np.argsort(kmeans_ind)]

corr_mds_kmeans_data.to_excel(os.path.join(period_corr_data_save_dir, prefix_labl+'_mds5_kmeans_corr_data_map.xlsx'))
save_path = os.path.join(period_corr_plot_save_dir, prefix_labl+'_mds5_kmeans_corr_data_map.png')
# plot(corr_mds_kmeans_data, save_path=save_path, figsize=(figsize)
plot_abs(corr_mds_kmeans_data, save_path=save_path, figsize=figsize)


